from tools.constant.redis_keys import ALGO_LOAD_MODEL, ALGO_LOAD_CONFIG


def hcytech_algo_task(load_model_func=None,release_model_func=None,reload_config_func=None):
    def decorator(func):
        def wrapper(*args, **kwargs):
            config = args[0]
            load_mode = config[ALGO_LOAD_MODEL]
            update_config = config[ALGO_LOAD_CONFIG]
            res = {
                "code":200,
                "message":"ok",
                "data":None
            }
            print(kwargs)
            try:
                # 前置处理
                if load_mode and load_model_func:
                    print("jiasdfsdf")
                    load_model_func()
                # 执行原始函数
                result = func(*args, **kwargs)
                res["data"] = result
                # 后置处理
                if not load_mode and release_model_func:
                    release_model_func()
                if update_config and reload_config_func:
                    reload_config_func()
            except Exception as e:
                res = {
                    "code": 500,
                    "message": "算法调用异常",
                    "data": str(e)
                }
            return res
        return wrapper
    return decorator